#ifndef BALL_TRACKER_H
#define BALL_TRACKER_H

#include "tracknetv3_trt.h"
#include "video_processor.h"
#include "yuv2rgb.h"
#include <opencv2/opencv.hpp>
#include <vector>
#include <string>
#include <fstream>

struct BallPosition {
    int x;
    int y;
    float confidence;
    int frame_id;
};

class BallTracker {
public:
    BallTracker();
    ~BallTracker();
    
    std::string getCSVPath() const { return csv_output_path_; }
    std::string getVideoPath() const { return video_output_path_; }
    int getFrameCount() const { return frame_count_; }
    
    bool init(const std::string& engine_path);
    bool initVideoWriter(const std::string& video_path, int width, int height, int fps);
    bool processFrame(const cv::Mat& frame);
    bool processFrameGPU(VideoProcessor& video_processor); // 新增GPU处理接口
    
    bool processFrameSoftware(VideoProcessor& video_processor);

    bool initGPUBuffers(int width, int height);

    const BallPosition& getCurrentPosition() const { return current_position_; }
    const std::vector<BallPosition>& getTrackHistory() const { return track_history_; }
    
    void setOutputPaths(const std::string& video_path, const std::string& csv_path);
    bool writeOutputs();
    
    void drawTrackingResult(cv::Mat& frame);
    
private:
    TrackNetV3_TRT* tracknet_;
    std::vector<cv::Mat> frame_buffer_;
    BallPosition current_position_;
    std::vector<BallPosition> track_history_;
    
    // GPU处理相关
    unsigned char* d_y_plane_;
    unsigned char* d_uv_plane_;
    float* d_rgb_buffer_;
    float* h_rgb_buffer_;
    
    cv::VideoWriter video_writer_;
    std::ofstream csv_file_;
    std::string video_output_path_;
    std::string csv_output_path_;
    
    int input_width_;
    int input_height_;
    int input_channels_;
    int video_width_;
    int video_height_;
    float ratio_x_;
    float ratio_y_;
    int frame_count_;
    
    void cleanupGPUBuffers();
};

#endif // BALL_TRACKER_H
